[mlir][tensor] Preserve encoding when folding empty#176427
Merged
sommerlukas merged 3 commits intollvm:mainfrom Jan 19, 2026
Merged
[mlir][tensor] Preserve encoding when folding empty#176427sommerlukas merged 3 commits intollvm:mainfrom
sommerlukas merged 3 commits intollvm:mainfrom
Conversation
Addresses a long-live TODO to not drop the encoding when folding a `tensor.empty` with a reshape operation (`tensor.expand_shape`, `tensor.collapse_shape`). Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
Member
|
@llvm/pr-subscribers-mlir-tensor @llvm/pr-subscribers-mlir Author: Lukas Sommer (sommerlukas) ChangesAddresses a long-standing TODO to not drop the encoding when folding a Full diff: https://github.com/llvm/llvm-project/pull/176427.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
index 670865de6031f..b88c2886095a4 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
@@ -40,11 +40,14 @@ struct FoldEmptyTensorWithReshapeOp : public OpRewritePattern<ReshapeOp> {
!llvm::hasSingleElement(resultShapes))
return failure();
+ Attribute encoding;
+ if(auto tensorTy = dyn_cast<RankedTensorType>(reshapeOp.getResultType())){
+ encoding = tensorTy.getEncoding();
+ }
// Create new tensor.empty op.
- // TODO: Do not drop tensor type encoding.
Value emptyTensor =
EmptyOp::create(rewriter, loc, resultShapes[0],
- reshapeOp.getResultType().getElementType());
+ reshapeOp.getResultType().getElementType(), encoding);
if (emptyTensor.getType() != reshapeOp.getResultType()) {
rewriter.replaceOpWithNewOp<tensor::CastOp>(
reshapeOp, reshapeOp.getResultType(), emptyTensor);
diff --git a/mlir/test/Dialect/Tensor/fold-empty-op.mlir b/mlir/test/Dialect/Tensor/fold-empty-op.mlir
index 7b11c9f43c7ec..62ee7e8c2d5ca 100644
--- a/mlir/test/Dialect/Tensor/fold-empty-op.mlir
+++ b/mlir/test/Dialect/Tensor/fold-empty-op.mlir
@@ -37,6 +37,27 @@ func.func @empty_reshape_collapse(%arg0 : index) -> tensor<6x5x?xf32> {
// CHECK-NEXT: %[[INIT:.+]] = tensor.empty(%[[D]])
// CHECK-NEXT: return %[[INIT]]
+#encoding = #test.tensor_encoding<"encoding">
+
+func.func @empty_expand_encoding() -> tensor<2x3x4x2xf32, #encoding> {
+ %0 = tensor.empty() : tensor<6x8xf32, #encoding>
+ %1 = tensor.expand_shape %0 [[0, 1], [2, 3]] output_shape [2, 3, 4, 2] : tensor<6x8xf32, #encoding> into tensor<2x3x4x2xf32, #encoding>
+ return %1 : tensor<2x3x4x2xf32, #encoding>
+}
+// CHECK-LABEL: func.func @empty_expand_encoding
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x3x4x2xf32, #test.tensor_encoding<"encoding">>
+// CHECK-NEXT: return %[[INIT]]
+
+func.func @empty_collapse_encoding() -> tensor<6x8xf32, #encoding> {
+ %0 = tensor.empty() : tensor<2x3x4x2xf32, #encoding>
+ %1 = tensor.collapse_shape %0 [[0, 1], [2, 3]]
+ : tensor<2x3x4x2xf32, #encoding> into tensor<6x8xf32, #encoding>
+ return %1 : tensor<6x8xf32, #encoding>
+}
+// CHECK-LABEL: func.func @empty_collapse_encoding
+// CHECK: %[[EMPTY_0:.*]] = tensor.empty() : tensor<6x8xf32, #test.tensor_encoding<"encoding">>
+// CHECK-NEXT: return %[[EMPTY_0]]
+
func.func @fold_empty_tensor_with_slice
(%arg0 : index, %arg1 : index) -> tensor<5x?x20xf32>
{
|
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
BStott6
pushed a commit
to BStott6/llvm-project
that referenced
this pull request
Jan 22, 2026
Addresses a long-standing TODO to not drop the encoding when folding a `tensor.empty` with a reshape operation (`tensor.expand_shape`, `tensor.collapse_shape`). --------- Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
amd-eochoalo
added a commit
to iree-org/iree
that referenced
this pull request
Jan 22, 2026
Reverts carried forward: Local revert of llvm/llvm-project#169614 due to #22649 * Also integrates torch-mlir * Fixes test due to llvm/llvm-project#176427
keshavvinayak01
pushed a commit
to iree-org/iree
that referenced
this pull request
Jan 27, 2026
Reverts carried forward: Local revert of llvm/llvm-project#169614 due to #22649 * Also integrates torch-mlir * Fixes test due to llvm/llvm-project#176427 Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Addresses a long-standing TODO to not drop the encoding when folding a
tensor.emptywith a reshape operation (tensor.expand_shape,tensor.collapse_shape).